import data_util
import model
import mindspore as ms
from tqdm import tqdm
import mindspore.ops as ops

def test(args_opt,bert_config,network=None):
    if args_opt.split_or_not==1:
        split=True
    else:
        split=False
    # load data
    raw=data_util.load_testset(args_opt.test_data_path)
    print('load data successfully')

    # data preprocessing
    test_dataset=data_util.pre_process(raw,args_opt.vocab_path,args_opt.max_seq_len,args_opt.batch_size,False)
    print('pre-process data successfully')

    # construct model
    if network==None:
        net = model.classifier_predict(bert_config,False,args_opt.fc_dim1,args_opt.class_num,args_opt.model_weight)
        ms.load_checkpoint(args_opt.model_weight,net)
        print('load model successully')
        net.set_train(False)
    else:
        if split==False:
            net=network
            net.set_train(False)
        else:
            net=network
            ms.load_checkpoint(args_opt.out_path+"/"+"best_acc.ckpt",net)
            net.set_train(False)

    #information
    print("============== Starting testing ==============")
    print("one epoch: %i steps"%test_dataset.get_dataset_size())

    argmax_ops=ops.Argmax()

    #predict
    results=[]
    total = test_dataset.get_dataset_size()
    with tqdm(total=total) as t:
        for i in test_dataset.create_tuple_iterator():
            t.set_description('testing')
            predictions = net(i[0],i[1],i[2])

            result=argmax_ops(predictions)
            result=result+1
            #print(result)
            result_=[int(j) for j in result]
            results=results+result_
            t.update(1)
    
    print('训练完成，结果写回')
    f=open(args_opt.result_path+'bert_test.txt','w')
    for number in results:
        f.write(str(number)+'\n')
    f.close()

def test_tiny_model(args_opt,network):
    # load data
    raw=data_util.load_testset(args_opt.test_data_path)
    print('load data successfully')

    # data preprocessing
    test_dataset=data_util.pre_process(raw,args_opt.vocab_path,args_opt.max_seq_len,args_opt.batch_size,False)
    print('pre-process data successfully')

    # construct model
    net=network
    net.set_train(False)

    #information
    print("============== Starting testing ==============")
    print("one epoch: %i steps"%test_dataset.get_dataset_size())

    argmax_ops=ops.Argmax()

    #predict
    results=[]
    total = test_dataset.get_dataset_size()
    with tqdm(total=total) as t:
        for i in test_dataset.create_tuple_iterator():
            t.set_description('testing')
            predictions = net(i[0])

            result=argmax_ops(predictions)
            result=result+1
            result_=[int(j) for j in result]
            results=results+result_
            t.update(1)
    
    print('训练完成，结果写回')
    f=open(args_opt.result_path+'bert_test.txt','w')
    for number in results:
        f.write(str(number)+'\n')
    f.close()